package com.wangbo.heap;

import java.util.*;

/**
 * 自定义堆：
 * 增加 resign() 方法，允许用户然后手动校准
 */
public class MyHeap<T> {
    public int heapSize;
    public List<T> heap;
    public Map<T, Integer> indexMap; // 元素坐标映射
    public Comparator<T> comparator;

    public MyHeap(Comparator<T> comparator) {
        this.heapSize = 0;
        this.heap = new ArrayList<>();
        indexMap = new HashMap<>();
        this.comparator = comparator;
    }

    public boolean isEmpty() {
        return heapSize == 0;
    }

    public int size() {
        return heapSize;
    }

    public boolean contain(T key) {
        return indexMap.containsKey(key);
    }

    /**
     * 添加新节点
     */
    public void push(T value) {
        if (contain(value)) {
            return;
        }

        int index = heapSize++;
        indexMap.put(value, index);
        heap.add(value);
        heapInsert(index);
    }

    /**
     * 堆调整(向上)
     * index节点和左右孩子中的较大方比较大小，index节点较小则交换位置，继续向下比较，直到index节点大于等于或者没有孩子
     */
    public void heapInsert(int index) {
        int parentIndex = (index - 1) / 2;
        while (comparator.compare(heap.get(index), heap.get(parentIndex)) < 0) {
            swap(index, parentIndex);
            index = parentIndex;
            parentIndex = (index - 1) / 2;
        }
    }

    /**
     * 交换元素位置
     */
    public void swap(int i, int j) {
        T obj1 = heap.get(i);
        T obj2 = heap.get(j);

        // 交换元素位置
        heap.set(i, obj2);
        heap.set(j, obj1);

        // 更新元素位置坐标
        indexMap.put(obj1, j);
        indexMap.put(obj2, i);
    }

    /**
     * 弹出根节点
     */
    public T pop() {
        T result = heap.get(0);
        int tail = heapSize - 1;
        swap(0, tail);
        heap.remove(tail);
        indexMap.remove(result);
        heapify(0, --heapSize);
        return result;
    }

    /**
     * 堆调整(向下)
     */
    public void heapify(int index, int heapSize) {
        int left = index * 2 + 1;
        while (left < heapSize) {
            int right = left + 1;
            int bigger = right < heapSize && (comparator.compare(heap.get(left), heap.get(right)) < 0) ? right : left;
            bigger = (comparator.compare(heap.get(index), heap.get(bigger)) < 0) ? bigger : index; // index和largest的较大方
            if (bigger == index) {
                break;
            } else {
                swap(bigger, index);
                index = bigger;
                left = index * 2 + 1;
            }
        }
    }

    /**
     * 手动校准堆
     */
    public void resign(T value) {
        int index = indexMap.get(value);
        heapInsert(index);
        heapify(index, heapSize);
    }
}
